import os
import numpy as np
import librosa
import soundfile as sf

# 全局变量
PROCESSED_DATA_PATH = '../data/raw/processed'  # 原始音频文件路径
AUGMENTED_DATA_PATH = '../data/raw/augmented'  # 增强后的音频文件路径

# 创建增强数据保存路径
if not os.path.exists(AUGMENTED_DATA_PATH):
    os.makedirs(AUGMENTED_DATA_PATH)


def time_stretch(data, rate=1.0):
    """
    改变音频的播放速度
    :param data: 输入音频数据
    :param rate: 播放速度调整比例
    :return: 拉伸后的音频数据
    """
    return librosa.effects.time_stretch(data, rate)


def pitch_shift(data, sr, n_steps):
    """
    改变音频的音高
    :param data: 输入音频数据
    :param sr: 采样率
    :param n_steps: 调整音高的步数（正数提高音高，负数降低音高）
    :return: 调整音高后的音频数据
    """
    return librosa.effects.pitch_shift(data, sr, n_steps)


def load_and_augment_data():
    """
    加载和增强音频数据
    """
    for file in os.listdir(PROCESSED_DATA_PATH):
        if file.endswith('.npy'):
            # 加载音频数据
            file_path = os.path.join(PROCESSED_DATA_PATH, file)
            audio_data = np.load(file_path)
            sr = 22050  # 假设采样率为 22050 Hz

            # 进行音频拉伸
            stretched_data_fast = time_stretch(audio_data, 1.25)
            stretched_data_slow = time_stretch(audio_data, 0.8)

            # 进行音调调整
            pitch_shifted_up = pitch_shift(audio_data, sr, 2)
            pitch_shifted_down = pitch_shift(audio_data, sr, -2)

            # 保存增强后的音频数据
            base_filename = os.path.splitext(file)[0]
            sf.write(os.path.join(AUGMENTED_DATA_PATH, base_filename + '_fast_1.25.wav'), stretched_data_fast, sr)
            sf.write(os.path.join(AUGMENTED_DATA_PATH, base_filename + '_slow_0.8.wav'), stretched_data_slow, sr)
            sf.write(os.path.join(AUGMENTED_DATA_PATH, base_filename + '_pitch_up_2.wav'), pitch_shifted_up, sr)
            sf.write(os.path.join(AUGMENTED_DATA_PATH, base_filename + '_pitch_down_2.wav'), pitch_shifted_down, sr)


if __name__ == '__main__':
    load_and_augment_data()
